package main

import (
	"bytes"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"sniffer/nbt"
	"sniffer/protocol"
	"strconv"
	"strings"
	"sync"
	"time"
)

type IPRange struct {
	IPs []net.IP
	mu  sync.Mutex
}

func (ips *IPRange) GetLen() int {
	ips.mu.Lock()
	length := len(ips.IPs)
	ips.mu.Unlock()

	return length
}

func (ips *IPRange) PopIP() net.IP {
	ips.mu.Lock()
	result := ips.IPs[0]
	ips.IPs = ips.IPs[1:]
	ips.mu.Unlock()

	return result
}

func GetIpRange(rangeString string) *IPRange {
	result := &IPRange{}
	IPs := make([]net.IP, 0)
	ipBoundaries := strings.Split(rangeString, "-")
	if len(ipBoundaries) == 1 {
		IPs = append(IPs, net.ParseIP(ipBoundaries[0]).To4())
	} else {
		ipStart := net.ParseIP(ipBoundaries[0]).To4()
		ipEnd := net.ParseIP(ipBoundaries[1]).To4()

		for ip := binary.BigEndian.Uint32(ipStart); ip <= binary.BigEndian.Uint32(ipEnd); ip++ {
			IPs = append(IPs, net.IPv4(byte(ip>>24), byte(ip>>16), byte(ip>>8), byte(ip)).To4())
		}
	}

	result.IPs = IPs

	return result
}

func GetPortRange(rangeString string) []uint16 {
	ports := make([]uint16, 0)

	portBoundaries := strings.Split(rangeString, "-")
	if len(portBoundaries) == 1 {
		portStart, err := strconv.ParseUint(portBoundaries[0], 10, 16)
		if err != nil {
			return ports
		}
		ports = append(ports, uint16(portStart))
	} else {
		portStart, err := strconv.ParseUint(portBoundaries[0], 10, 16)
		if err != nil {
			return ports
		}
		portEnd, err := strconv.ParseUint(portBoundaries[1], 10, 16)
		if err != nil {
			return ports
		}

		for port := uint16(portStart); port <= uint16(portEnd); port++ {
			ports = append(ports, port)
		}
	}

	return ports
}

func worker(wg *sync.WaitGroup, ips *IPRange, ports []uint16, foundServers *nbt.McServers, timeout uint) {
	log.Println("Starting worker")

	for {
		if ips.GetLen() == 0 {
			wg.Done()
			log.Println("Worker done")
			return
		}

		ip := ips.PopIP()
		for _, port := range ports {
			conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", ip, port), time.Second*time.Duration(timeout))
			if err != nil {
				log.Printf("Failed to connect: %s\n", err)
				continue
			}

			hs, _ := protocol.NewHandshakePacket(765, ip.String(), uint16(port))
			s, _ := protocol.NewStatusPacket()

			if _, err := hs.WriteTo(conn); err != nil {
				log.Printf("Failed to write handshake: %s\n", err)
				continue
			}
			if _, err := s.WriteTo(conn); err != nil {
				log.Printf("Failed to write status: %s\n", err)
				continue
			}

			resp, err := protocol.ReadPacket(conn)
			if err != nil {
				log.Printf("Failed to read response: %s\n", err)
				continue
			}

			data, err := resp.ToMcStatusResponse()
			if err != nil {
				log.Printf("Failed to convert response to data: %s\n", err)
				continue
			}

			log.Printf("Found a server: %s:%d %s %d/%d\n", ip, port, data.Version.Name, data.Players.Online, data.Players.Max)

			var server nbt.McServer

			server.Ip = fmt.Sprintf("%s:%d", ip, port)
			server.Name = fmt.Sprintf("#%d %s - %s:%d", foundServers.GetLen(), data.Version.Name, ip, port)

			foundServers.AddServer(server)
		}
	}

}

func main() {
	workerCount := flag.Uint("workers", 10, "Number of worker goroutines")
	timeout := flag.Uint("timeout", 7, "Timeout in seconds before aborting connection")
	ipRange := flag.String("ip-range", "", "Server IP range")
	portRange := flag.String("port-range", "25565", "Server port range")
	outputFile := flag.String("output-file", "", "NBT filename")

	flag.Parse()
	if len(*ipRange) == 0 || len(*portRange) == 0 {
		flag.PrintDefaults()
		os.Exit(-1)
	}

	ipQueue := GetIpRange(*ipRange)
	ports := GetPortRange(*portRange)
	foundServers := nbt.NewMcServers()

	var wg sync.WaitGroup
	log.Printf("Starting %d workers\n", *workerCount)

	for i := uint(0); i < *workerCount; i++ {
		wg.Add(1)
		go worker(&wg, ipQueue, ports, foundServers, *timeout)
	}

	wg.Wait()

	fmt.Print("\n\n\n")
	fmt.Printf("=========== FOUND %d SERVERS ===========\n", len(foundServers.Servers))

	for _, server := range foundServers.Servers {
		fmt.Printf("%s - %s\n", server.Name, server.Ip)
	}

	if len(*outputFile) != 0 {
		var buffer bytes.Buffer
		if _, err := foundServers.WriteTo(&buffer); err != nil {
			fmt.Printf("Failed to serialize found servers: %s\n", err)
			os.Exit(-1)
		}

		file, err := os.Create(*outputFile)
		if err != nil {
			fmt.Printf("Failed to open file for writing: %s\n", err)
			os.Exit(-1)
		}

		n, err := io.Copy(file, &buffer)
		if err != nil {
			fmt.Printf("Failed to copy data from buffer to output file: %s\n", err)
			os.Exit(-1)
		}

		fmt.Printf("Saved servers to %s, written %d bytes\n", *outputFile, n)
	}
}
